import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

def ssr_channel_raw(channel, sigma):
    channel = channel.astype(np.float32) / 255.0
    illumination = gaussian_filter(channel, sigma=sigma)
    with np.errstate(divide='ignore', invalid='ignore'):
        retinex = np.log(channel+0.5) - np.log(illumination+0.5)
        retinex[np.isneginf(retinex)] = np.min(retinex[np.isfinite(retinex)])
    return retinex

def msr_channel_weighted(channel, sigmas, weights):
    # 权重和sigmas长度相同，权重和为1
    retinex_scales = [w * ssr_channel_raw(channel, s) for s, w in zip(sigmas, weights)]
    retinex = sum(retinex_scales)
    return retinex

def msrcr_image(
    img_pil,
    sigmas=[30, 80, 100],
    weights=[0.33, 0.33, 0.34],  # MSR权重固定和为1
    alpha=125.0,
    c=1.0,
    gain=1.0,
    bias=0.0,     # 颜色恢复函数后的偏置项
):
    img = np.array(img_pil).astype(np.float32) / 255.0

    channels = []
    denom = np.sum(img, axis=2) + 1e-6  # 防止除零

    for i in range(3):
        msr = msr_channel_weighted(img[..., i], sigmas, weights)
        with np.errstate(divide='ignore', invalid='ignore'):
            crf = np.log(alpha * (img[..., i] / denom) + c)
        msrcr = gain * msr * crf + bias
        channels.append(msrcr)

    result = np.stack(channels, axis=-1)
    result -= np.min(result)
    result /= np.max(result)
    result = (result * 255).astype(np.uint8)

    return Image.fromarray(result)

# 示例调用和显示
if __name__ == "__main__":
    img = Image.open("D:\Retinex\\test\\test.jpg").convert("RGB")
    msrcr = msrcr_image(img)
    plt.figure(figsize=(8, 4))
    plt.subplot(1, 2, 1)
    plt.imshow(img)
    plt.title("Original")
    plt.axis("off")

    plt.subplot(1, 2, 2)
    plt.imshow(msrcr)
    plt.title("MSRCR with gain & bias")
    plt.axis("off")
    plt.tight_layout()
    plt.show()
